import numpy as np
import pyopencl as cl
from verbose import Verbose
echo = Verbose().echo

#mako
N_TYPE = 1
W_P_PAR = 4   #
test_p_par = np.zeros((N_TYPE,N_TYPE,W_P_PAR),'float32')
test_p_par[0,0] = [1., 1.*1., 2.5*2.5, 0]

class Pair:
    """
    ---
    ,can be const in cl file by mako
    float4: p_par[ntype][ntype][*] #pair_parameters
    int*: neigh_pair,neigh_bond
    =====
    """
        
    def __init__(self,name,cl_file,p_par,context,flags=''):
        self.name = name #'pair_lj_cut'
        self.cl_file = cl_file #'pair_lj_cut.cl'
        self.p_par = p_par
        #presetting parameters:
        self.p_par = test_p_par
        #copy parameters:
        mf = cl.mem_flags
        self._p_par = cl.Buffer(context,mf.READ_ONLY|mf.COPY_HOST_PTR,hostbuf=self.p_par)
        #compile opencl source:
        cl_code = file(self.cl_file,'r').read()
        self.prg = cl.Program(context,cl_code).build(options=flags)
        echo('-----pair-%s parameters-----'%(self.name),12) #debug
        echo(self.p_par,12) #debug
        echo('---end pair-%s parameters---'%(self.name),12) #debug

    def read_parameter(self,parfile):
        """un-finished """
        self.box_p = np.array([1,1,1,0],'float32')
        keywords = [ "pair_style", "pair_coeff" ]
        #f = file(parfile,'r')
        nt = self.n_atom_type
        self.pair_type = {} #for more pair types, need mako !!! 
        p_par = np.ndarray((nt,nt,W_P_PAR),'float32')  #parameter index
        for line in file(parfile,'r'):
            s = line.strip().split()
            if   s[0] == "pair_style":
                pass
            elif s[0] == "pair_coeff":
                i = int(s[1]) - 1 #lammps index from 1 not 0
                j = int(s[2]) - 1
                p_par[i,j,0] = float(s[3])
                p_par[i,j,1] = float(s[4])
                p_par[i,j,2] = float(s[5])
        #f.readline()
        #f.close()
        self.p_par = p_par
        echo('read parameter file- %s: done.'%(parfile),10) #debug

    def add_f_e_slow(self,que,sglb,sloc,cfg): 
        self.prg.add_f_e_slow(que,sglb,sloc,
                np.int32(cfg.n_atoms),cfg.box_l,cfg.box_p,
                cfg._r_t,cfg._f_e,cfg._vir,self._p_par)

    def add_f_e_vir_all(self,que,sglb,sloc,cfg,nei): 
        self.prg.add_f_e_vir_all(que,sglb,sloc,
                np.int32(cfg.n_atoms),cfg.box_l,cfg.box_p,
                cfg._r_t,cfg._f_e,cfg._vir,self._p_par,nei._nei,nei._n_nei)
